package com.gjy.service.es.util;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.search.aggregations.Aggregation;
import org.elasticsearch.search.aggregations.AggregationBuilders;
import org.elasticsearch.search.aggregations.Aggregations;
import org.elasticsearch.search.aggregations.bucket.terms.ParsedStringTerms;
import org.elasticsearch.search.aggregations.bucket.terms.Terms;
import org.elasticsearch.search.aggregations.bucket.terms.TermsAggregationBuilder;
import org.springframework.data.elasticsearch.client.ClientConfiguration;
import org.springframework.data.elasticsearch.client.RestClients;
import org.springframework.util.ObjectUtils;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.atomic.AtomicReference;

/**
 * @author 宫静雨
 * @version 1.0
 * @since 2023-07-24 09:13:08
 */
public class Util {

    public static RestHighLevelClient client() {
        String uri = "192.168.200.152:9200";
        ClientConfiguration build = ClientConfiguration.builder().connectedTo(uri).build();
        return RestClients.create(build).rest();
    }

    public static List<TermsAggregationBuilder> termsParallel(List<String> terms) {
        List<TermsAggregationBuilder> list = Lists.newArrayList();
        for (String term : terms) {
            TermsAggregationBuilder builder = AggregationBuilders.terms(term).field(term);
            list.add(builder);
        }
        return list;
    }

    public static Map<String, Map<String, Object>> termsParallelResult(List<String> terms, Aggregations aggregations) {
        Map<String, Map<String, Object>> ms = new HashMap<>();
        for (String term : terms) {
            List<? extends Terms.Bucket> buckets = ((ParsedStringTerms) aggregations.get(term)).getBuckets();
            HashMap<String, Object> map = new HashMap<>();
            for (Terms.Bucket bucket : buckets) {
                String key = bucket.getKeyAsString();
                long value = bucket.getDocCount();
                map.put(key, value);
            }
            ms.put(term, map);
        }
        return ms;
    }

    public static TermsAggregationBuilder termsNested(List<String> terms) {
        AtomicReference<TermsAggregationBuilder> root = new AtomicReference<>();
        AtomicReference<TermsAggregationBuilder> cursor = new AtomicReference<>();

        terms.forEach(term -> {
            TermsAggregationBuilder builder = AggregationBuilders.terms(term).field(term);
            if (Objects.isNull(root.get())) {
                root.set(builder);
                cursor.set(root.get());
            } else {
                cursor.get().subAggregation(builder);
                cursor.set(builder);
            }
        });
        return root.get();
    }

    public static Map<String, Object> termsNestedResult(List<String> terms, int index, Aggregations aggregations) {
        Map<String, Object> map = Maps.newHashMap();

        Aggregation aggregation = aggregations.get(terms.get(index));
        List<? extends Terms.Bucket> buckets = ((ParsedStringTerms) aggregation).getBuckets();

        for (Terms.Bucket bucket : buckets) {
            String key = bucket.getKeyAsString();
            long value = bucket.getDocCount();
            map.put(key, value);

            if (index < terms.size() - 1) {
                Aggregations as1 = bucket.getAggregations();
                Aggregation a1 = as1.get(terms.get(index + 1));
                if (!ObjectUtils.isEmpty(as1) && !ObjectUtils.isEmpty(a1)) {
                    Map<String, Object> m1 = termsNestedResult(terms, index + 1, as1);
                    m1.forEach((k, v) -> map.put(key + "_" + k, v));
                }
            }
        }
        return map;
    }

}
